package org.example.myleet.p814;

import org.example.myleet.Utils.TreeNode;

public class Solution {
    public TreeNode pruneTree(TreeNode root) {
        pruneDfs(root);
        if (0 == root.val && null == root.left && null == root.right) {
            return null;
        }
        return root;
    }

    private int pruneDfs(TreeNode node) {
        if (null == node.left && null == node.right) {
            return node.val;
        }
        int vLeft = 0;
        if (null != node.left) {
            vLeft = pruneDfs(node.left);
            if (0 == vLeft) {
                node.left = null;
            }
        }
        int vRight = 0;
        if (null != node.right) {
            vRight = pruneDfs(node.right);
            if (0 == vRight) {
                node.right = null;
            }
        }
        int vSum = vLeft + vRight + node.val;
        return vSum > 0 ? 1 : 0;
    }
}
